import numpy as np
from metrics import r2_score


class LinearRegression:

	def __init__(self):
		self.coef_ = None
		self.intercept_ = None
		self._theta = None

	def fit_normal(self, X_train, y_train):
		assert X_train.shape[0] == y_train.shape[0], "size must be same"

		X_b = np.hstack([np.ones((X_train.shape[0], 1)), X_train])
		self._theta = np.linalg.inv(X_b.T.dot(X_b)).dot(X_b.T).dot(y_train)
		self.coef_ = self._theta[1:]
		self.intercept_ = self._theta[0]

		return self

	def fit_gd(self, X_train, y_train, eta=0.01, n_iters=1e4):
		assert X_train.shape[0] == y_train.shape[0], "size must be same"

		def J(theta, X_b, y):
			try:
				return np.sum((y - X_b.dot(theta)) ** 2) / len(X_b)
			except:
				return float('inf')

		def dJ(theta, X_b, y):
			# res = np.empty(len(theta))
			# res[0] = np.sum(X_b.dot(theta) - y)
			#
			# for i in range(1, len(theta)):
			# 	res[i] = np.sum((X_b.dot(theta) - y).dot(X_b[:, i]))
			#
			# return res * 2 / len(y)

			return X_b.T.dot(X_b.dot(theta) - y) * 2 / len(y)

		def gradient_descent(X_b, y, initial_theta, eta, n_iters=1e4, epsilon=1e-8):
			theta = initial_theta
			cur_iter = 0

			while cur_iter < n_iters:
				gradient = dJ(theta, X_b, y)
				last_theta = theta
				theta = theta - eta * gradient

				if abs(J(theta, X_b, y) - J(last_theta, X_b, y)) < epsilon:
					break

				cur_iter += 1

			return theta

		X_b = np.hstack([np.ones((len(X_train), 1)), X_train])
		initial_theta = np.zeros(X_b.shape[1])

		self._theta = gradient_descent(X_b, y_train, initial_theta, eta, n_iters)

		self.coef_ = self._theta[1:]
		self.intercept_ = self._theta[0]

		return self

	def fit_sgd(self, X_train, y_train, n_iters = 5, t0 = 5, t1 = 50):
		assert X_train.shape[0] == y_train.shape[0], "size must be same"
		assert n_iters >= 1

		def dJ_sgd(X_b_i, y_i, theta):
			return X_b_i.dot(X_b_i.dot(theta) - y_i) * 2

		def sgd(X_b, y, initial_theta, n_iters, t0 = 5, t1= 50):

			def learning_rate(t):
				return t0 / (t + t1)

			theta = initial_theta
			m = len(X_b)

			for cur_iter in range(n_iters):
				indexes = np.random.permutation(m)
				X_b_new = X_b[indexes]
				y_new = y[indexes]

				for i in range(m):
					gradient = dJ_sgd(X_b_new[i], y_new[i], theta)
					theta = theta - learning_rate(cur_iter * m + i) * gradient

				cur_iter += 1

			return theta

		X_b = np.hstack([np.ones((len(X_train), 1)), X_train])
		initial_theta = np.zeros(X_b.shape[1])

		self._theta = sgd(X_b, y_train, initial_theta, n_iters)

		self.coef_ = self._theta[1:]
		self.intercept_ = self._theta[0]

		return self

	def predict(self, X_test):
		assert self.coef_ is not None and self.intercept_ is not None, "must fit before predict"
		assert X_test.shape[1] == len(self.coef_), "feature number must be same"

		X_b = np.hstack([np.ones((X_test.shape[0], 1)), X_test])
		return X_b.dot(self._theta)

	def score(self, X_test, y_test):
		assert X_test.shape[0] == len(y_test), "length must be same"

		return r2_score(y_test, self.predict(X_test))

	def __repr__(self):
		return "LinearRegression()"
